package helper.chart

import helper.base.debug.DebugHelper.DebuggerList
import helper.base.debug.DebugHelper.DebuggerList.Companion.makeDebuggerList
import helper.chart.ChartHelper.JFreeChartEx.taskDisplayInFrame
import helper.chart.ChartHelper.JFreeChartEx.taskSaveToFile
import org.apache.commons.math3.distribution.PoissonDistribution
import org.jfree.chart.ChartFactory
import org.jfree.chart.JFreeChart
import org.jfree.chart.plot.PlotOrientation
import org.jfree.data.category.DefaultCategoryDataset

object PoissonChartHelperKt {
    object PoissonChartHelper {
        fun getPoissonDistributionChart(xMax: Int, lambda: Double): JFreeChart {
            val x = IntArray(xMax) { it }
            val poisson = PoissonDistribution(lambda)
            val pmf = DoubleArray(xMax) { poisson.probability(it) }

            return DefaultCategoryDataset().apply {
                for (i in x.indices) {
                    addValue(pmf[i], "Poisson PMF", x[i])
                }
            }.let {
                ChartFactory.createBarChart(
                    "Poisson PMF with λ=$lambda",
                    "Number of Events",
                    "Probability",
                    it,
                    PlotOrientation.VERTICAL,
                    true,
                    true,
                    false
                )
            }
        }

        fun taskPoissonDistributionChart(
            // 设置参数 lambda
            lambda: Double,
            // 设置横坐标最大值
            xMax: Int,
            savedFileName: String? = null,
            showedFrameTitle: String? = null,
            debuggerList: DebuggerList = makeDebuggerList(),
        ) {
            val chart: JFreeChart = getPoissonDistributionChart(
                xMax,
                lambda,
            )
            savedFileName?.let {
                chart.taskSaveToFile(it)
            }
            showedFrameTitle?.let {
                chart.taskDisplayInFrame(it)
            }
        }
    }
}